/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package eu.marc.juliator.gui.peptideprotein;

import com.compomics.util.experiment.biology.Protein;
import com.compomics.util.experiment.identification.SequenceFactory;
import com.compomics.util.experiment.identification.matches.ProteinMatch;
import com.compomics.util.experiment.identification.protein_inference.proteintree.ProteinTree;
import eu.marc.juliator.gui.io.InputFile;
import java.io.IOException;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import javax.swing.table.DefaultTableModel;

/**
 * table model for the peptide to protein matching
 *
 * @author Marc
 */
public class ProteinTableModel extends DefaultTableModel {

    /**
     * the input file
     */
    private InputFile inputFile = null;
    /**
     * The number of lines
     */
    private int nLines = 0;
    /**
     * a line to peptide number map
     */
    private HashMap<Integer, Integer> rowToPeptideMap = new HashMap<Integer, Integer>();
    /**
     * peptide to first line map
     */
    private HashMap<Integer, Integer> peptideToRowMap = new HashMap<Integer, Integer>();
    /**
     * The column containing the peptide sequences
     */
    private int peptideColumn;
    /**
     * the number of amino-acids which should be looked for
     */
    private int nAA;
    /**
     * The peptide-protein tree
     */
    private ProteinTree proteinTree;
    /**
     * The sequence factory
     */
    private SequenceFactory sequenceFactory = SequenceFactory.getInstance();
    /**
     * The ms2 tolerance of the search
     */
    private Double ms2Tolerance;
    /**
     * The matching type
     */
    private ProteinMatch.MatchingType matchType;

    /**
     * Empty constructor for an empty table
     */
    public ProteinTableModel() {
    }

    /**
     * Constructor
     */
    public ProteinTableModel(InputFile inputFile, int peptideColumn, int nAA, ProteinTree proteinTree, ProteinMatch.MatchingType matchType, Double ms2Tolerance) {
        this.inputFile = inputFile;
        nLines = Math.min(inputFile.getNLines(), 1000);
        this.peptideColumn = peptideColumn;
        this.nAA = nAA;
        this.proteinTree = proteinTree;
        this.matchType = matchType;
        this.ms2Tolerance = ms2Tolerance;
    }

    @Override
    public int getRowCount() {
        if (inputFile == null || inputFile.getNLines() == 0) {
            return 0;
        }
        return nLines;
    }

    @Override
    public int getColumnCount() {
        if (inputFile == null) {
            return 0;
        }
        return inputFile.getNColumns() + 5;
    }

    @Override
    public String getColumnName(int column) {
        if (column < inputFile.getNColumns()) {
            try {
                return inputFile.getValueAt(0, column);
            } catch (Exception e) {
                e.printStackTrace();
                return "";
            }
        } else if (column == inputFile.getNColumns()) {
            return "Sequence";
        } else if (column == inputFile.getNColumns()+ 1) {
            return "Accession";
        } else if (column == inputFile.getNColumns() + 2) {
            return "AA before";
        } else if (column == inputFile.getNColumns() + 3) {
            return "AA after";
        } else if (column == inputFile.getNColumns() + 4) {
            return "terminal";
        }
        return "";
    }

    @Override
    public Object getValueAt(int row, int column) {
        try {
            Integer index = rowToPeptideMap.get(row);
            if (index == null) {
                index = peptideToRowMap.size();
                peptideToRowMap.put(index, row);
                String sequence = inputFile.getValueAt(index + 1, peptideColumn - 1);
                if (sequence != null && sequence.length() >= 3) {
                    int length = getMappingLength(sequence);
                    for (int i = 0; i < length; i++) {
                        rowToPeptideMap.put(row + i, index);
                        if (i > 0) {
                            nLines++;
                        }
                    }
                }
            }
            if (column < inputFile.getNColumns()) {
                return inputFile.getValueAt(index + 1, column);
            }
            String sequence = inputFile.getValueAt(index + 1, peptideColumn - 1);
            if (sequence != null && sequence.length() >= 3) {
                int length = getMappingLength(sequence);
                if (length > 0) {
                    if (column == inputFile.getNColumns()) {
                        int offset = row - peptideToRowMap.get(index);
                        return getSequence(sequence, offset);
                    } else if (column == inputFile.getNColumns() + 1) {
                        int offset = row - peptideToRowMap.get(index);
                        return getAccession(sequence, offset);
                    } else if (column == inputFile.getNColumns() + 2) {
                        int offset = row - peptideToRowMap.get(index);
                        return getAaBefore(sequence, offset);
                    } else if (column == inputFile.getNColumns() + 3) {
                        int offset = row - peptideToRowMap.get(index);
                        return getAaAfter(sequence, offset);
                    } else if (column == inputFile.getNColumns() + 4) {
                        int offset = row - peptideToRowMap.get(index);
                        return getTerminal(sequence, offset);
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return "";
    }

    @Override
    public Class getColumnClass(int columnIndex) {
        for (int i = 0; i < getRowCount(); i++) {
            if (getValueAt(i, columnIndex) != null) {
                return getValueAt(i, columnIndex).getClass();
            }
        }
        return String.class;
    }

    @Override
    public boolean isCellEditable(int rowIndex, int columnIndex) {
        return false;
    }

    @Override
    public void setValueAt(Object aValue, int rowIndex, int columnIndex) {
    }

    /**
     * Returns the number of rows needed for a given peptide sequence
     *
     * @param sequence the sequence of interest
     * @return the number of row needed
     * @throws IOException
     * @throws InterruptedException
     * @throws ClassNotFoundException
     */
    private int getMappingLength(String sequence) throws IOException, InterruptedException, ClassNotFoundException, SQLException {
        int result = 0;
        HashMap<String, HashMap<String, ArrayList<Integer>>> mapping = proteinTree.getProteinMapping(sequence, matchType, ms2Tolerance);
        for (String peptide : mapping.keySet()) {
            for (String protein : mapping.get(peptide).keySet()) {
                result += mapping.get(peptide).get(protein).size();
            }
        }
        return result;
    }

    /**
     * Returns the amino acid after the sequence at a given offset in the
     * protein mapping
     *
     * @param sequence the sequence of interest
     * @param offset the offset
     * @return the amino acid after
     * @throws IOException
     * @throws InterruptedException
     * @throws ClassNotFoundException
     */
    private String getAaAfter(String sequence, int offset) throws IOException, InterruptedException, ClassNotFoundException, SQLException {
        String result = "";
        HashMap<String, HashMap<String, ArrayList<Integer>>> mapping = proteinTree.getProteinMapping(sequence, matchType, ms2Tolerance);
        ArrayList<String> sequences = new ArrayList<String>(mapping.keySet());
        Collections.sort(sequences);
        int cpt = 0;
        for (String peptide : sequences) {
            ArrayList<String> proteins = new ArrayList<String>(mapping.get(peptide).keySet());
            Collections.sort(proteins);
            for (String accession : proteins) {
                ArrayList<Integer> indexes = mapping.get(peptide).get(accession);
                if (cpt + indexes.size() > offset) {
                    Collections.sort(indexes);
                    int aa = indexes.get(offset - cpt) + sequence.length();
                    Protein protein = sequenceFactory.getProtein(accession);
                    String proteinSequence = protein.getSequence();
                    for (int i = 0; i < nAA; i++) {
                        int index = aa + i;
                        if (index >= proteinSequence.length()) {
                            break;
                        }
                        result += proteinSequence.charAt(index);
                    }
                    break;
                }
                cpt += indexes.size();
            }
        }
        return result;
    }

    /**
     * Returns the amino acid before the sequence at a given offset in the
     * protein mapping
     *
     * @param sequence the sequence of interest
     * @param offset the offset
     * @return the amino acid after
     * @throws IOException
     * @throws InterruptedException
     * @throws ClassNotFoundException
     */
    private String getAaBefore(String sequence, int offset) throws IOException, InterruptedException, ClassNotFoundException, SQLException {
        String result = "";
        HashMap<String, HashMap<String, ArrayList<Integer>>> mapping = proteinTree.getProteinMapping(sequence, matchType, ms2Tolerance);
        ArrayList<String> sequences = new ArrayList<String>(mapping.keySet());
        Collections.sort(sequences);
        int cpt = 0;
        for (String peptide : sequences) {
            ArrayList<String> proteins = new ArrayList<String>(mapping.get(peptide).keySet());
            Collections.sort(proteins);
            for (String accession : proteins) {
                ArrayList<Integer> indexes = mapping.get(peptide).get(accession);
                if (cpt + indexes.size() > offset) {
                    Collections.sort(indexes);
                    int aa = indexes.get(offset - cpt) - nAA;
                    Protein protein = sequenceFactory.getProtein(accession);
                    String proteinSequence = protein.getSequence();
                    for (int i = 0; i < nAA; i++) {
                        int index = aa + i;
                        if (index < 0) {
                            break;
                        }
                        result += proteinSequence.charAt(index);
                    }
                    break;
                }
                cpt += indexes.size();
            }
        }
        return result;
    }

    /**
     * Returns the peptide sequence
     *
     * @param sequence the sequence of interest
     * @param offset the offset
     * @return the amino acid after
     * @throws IOException
     * @throws InterruptedException
     * @throws ClassNotFoundException
     */
    private String getSequence(String sequence, int offset) throws IOException, InterruptedException, ClassNotFoundException, SQLException {
        String result = "";
        HashMap<String, HashMap<String, ArrayList<Integer>>> mapping = proteinTree.getProteinMapping(sequence, matchType, ms2Tolerance);
        ArrayList<String> sequences = new ArrayList<String>(mapping.keySet());
        Collections.sort(sequences);
        int cpt = 0;
        for (String peptide : sequences) {
            ArrayList<String> proteins = new ArrayList<String>(mapping.get(peptide).keySet());
            Collections.sort(proteins);
            for (String accession : proteins) {
                ArrayList<Integer> indexes = mapping.get(peptide).get(accession);
                if (cpt + indexes.size() > offset) {
                    return peptide;
                }
                cpt += indexes.size();
            }
        }
        return result;
    }

    /**
     * Returns the protein accession
     *
     * @param sequence the sequence of interest
     * @param offset the offset
     * @return the amino acid after
     * @throws IOException
     * @throws InterruptedException
     * @throws ClassNotFoundException
     */
    private String getAccession(String sequence, int offset) throws IOException, InterruptedException, ClassNotFoundException, SQLException {
        String result = "";
        HashMap<String, HashMap<String, ArrayList<Integer>>> mapping = proteinTree.getProteinMapping(sequence, matchType, ms2Tolerance);
        ArrayList<String> sequences = new ArrayList<String>(mapping.keySet());
        Collections.sort(sequences);
        int cpt = 0;
        for (String peptide : sequences) {
            ArrayList<String> proteins = new ArrayList<String>(mapping.get(peptide).keySet());
            Collections.sort(proteins);
            for (String accession : proteins) {
                ArrayList<Integer> indexes = mapping.get(peptide).get(accession);
                if (cpt + indexes.size() > offset) {
                    return accession;
                }
                cpt += indexes.size();
            }
        }
        return result;
    }

    /**
     * Returns whether the given sequence is found at a terminus of the protein
     *
     * @param sequence the sequence of interest
     * @param offset the offset
     * @return the amino acid after
     * @throws IOException
     * @throws InterruptedException
     * @throws ClassNotFoundException
     */
    private String getTerminal(String sequence, int offset) throws IOException, InterruptedException, ClassNotFoundException, SQLException {
        String result = "";
        HashMap<String, HashMap<String, ArrayList<Integer>>> mapping = proteinTree.getProteinMapping(sequence, matchType, ms2Tolerance);
        ArrayList<String> sequences = new ArrayList<String>(mapping.keySet());
        Collections.sort(sequences);
        int cpt = 0;
        for (String peptide : sequences) {
            ArrayList<String> proteins = new ArrayList<String>(mapping.get(peptide).keySet());
            Collections.sort(proteins);
            for (String accession : proteins) {
                ArrayList<Integer> indexes = mapping.get(peptide).get(accession);
                if (cpt + indexes.size() > offset) {
                    Collections.sort(indexes);
                    Protein protein = sequenceFactory.getProtein(accession);
                    String proteinSequence = protein.getSequence();
                    if (proteinSequence.startsWith(sequence)) {
                        result += "N-term";
                    }
                    if (proteinSequence.endsWith(sequence)) {
                        if (!result.equals("")) {
                            result += " and ";
                        }
                        result += "C-term";
                    }
                    break;
                }
                cpt += indexes.size();
            }
        }
        return result;
    }
}
